import pandas as pd
from elasticsearch import Elasticsearch
from flask import Flask, jsonify, request
from milvus import Milvus

app = Flask(__name__)

cache = {}
word_database_df = pd.read_pickle("data/word_database.pkl")
for id, row in word_database_df.iterrows():
    word = row["word"]
    count = row["count"]
    vectors = [float(v) for v in row["vectors"].split(",")]
    cache[str(id)] = {"word": word, "count": count, "vectors": vectors}

index_content = open("static/index.html",encoding="utf8").read()

@app.route('/')
def index():
    return index_content

@app.route('/search_word')
def search_word():
    search = request.values.get("q")

    data = []
    if search:
        res = es.search(index='search_word_analyse', doc_type='_doc', body={
            "size": 30,
            "query": {
                "match": {
                    "word": search
                }
            }
        })

        hits = res.get('hits',{}).get('hits',{})
        if hits:
            data = [h.get('_source') for h in hits]
    return jsonify(data)


@app.route('/find_similar')
def find_similar():
    records = []
    vector = request.values.get("v")
    if vector:
        status, q_records = milvus.search(collection_name, top_k=50,
                           query_records=[[float(v) for v in vector.split(",")]],
                           params={'nprobe': 16})
        if status.code==0:
            for record_id in q_records.id_array[0]:
                record = cache[str(record_id)]
                records.append(record)
    return jsonify(records)


if __name__ == '__main__':
    es = Elasticsearch(["http://localhost:9200"])
    milvus = Milvus(host='localhost', port='19530')
    collection_name = 'search_word_analyse'
    app.run("0.0.0.0")

